import base64
import numpy as np
import hqyj_mqtt
import queue
import cv2


def b64_to_np(image):
    image_data = base64.b64decode(image['image'])
    image_np = np.frombuffer(image_data, dtype=np.uint8)
    img = cv2.imdecode(image_np, cv2.IMREAD_COLOR)
    return img


if __name__ == '__main__':
    q_mqtt_data = queue.Queue(5)
    i = 1
    mqtt_client = hqyj_mqtt.MQTTClient('127.0.0.1', 21883, 'xiaoyibb', 'xiaoyiaa', q_mqtt_data)
    while True:
        image = q_mqtt_data.get()
        if 'image' in image:
            img = b64_to_np(image)
            cv2.imshow('img', img)
            key = cv2.waitKey(1) & 0xFF
            if key == ord("q"):
                break
            elif key == ord('v'):
                cv2.imwrite(f'images/{i}.png', img)
                i += 1
                print('save successful')
            else:
                continue
